
import tensorrt as trt
from calibrator import DecoderCalibrator


onnx_file = "onnx/decoder_opt.onnx"
workspace = 20480
use_int8 = True

logger = trt.Logger(trt.Logger.VERBOSE)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

parser = trt.OnnxParser(network, logger)
success = parser.parse_from_file(onnx_file)
for idx in range(parser.num_errors):
    print(parser.get_error(idx))

if not success:
    pass  # Error handling code here

profile = builder.create_optimization_profile()
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace * 1 << 20)  # 1MB

config.flags = 1 << int(trt.BuilderFlag.FP16)
config.flags = config.flags & ~(1 << int(trt.BuilderFlag.TF32))
if use_int8:
    config.flags = config.flags | (1 << int(trt.BuilderFlag.INT8))
    config.int8_calibrator = DecoderCalibrator("decoder.cache", 20)

serialized_engine = builder.build_serialized_network(network, config)

with open("decoder_int8.trt", "wb") as f:
    f.write(serialized_engine)
